{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "882fb634-0efc-42bf-8cc5-79d6a7312e26",
   "metadata": {},
   "source": [
    "# `AzureOpenAI` API Examples\n",
    "\n",
    "You can also use an OpenAI model deployed into Azure AI.\n",
    "For this, you will provide a few pieces of information from the Azure AI playground:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6b887dab",
   "metadata": {
    "tags": [
     "parameters",
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "call_delay_secs = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "61665201-12a8-4588-bb54-801a367c7535",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# If using DefaultAzureCredential below\n",
    "from azure.identity import DefaultAzureCredential, get_bearer_token_provider\n",
    "\n",
    "# This is the name of the model deployed, such as 'gpt-4' or 'gpt-3.5-turbo\n",
    "model = os.getenv(\"AZUREAI_OPENAI_CHAT_MODEL\", \"Please set the model\")\n",
    "\n",
    "# This is the deployment URL, as provided in the Azure AI playground ('view code')\n",
    "# It will end with 'openai.azure.com'\n",
    "azure_endpoint = os.getenv(\"AZUREAI_OPENAI_CHAT_ENDPOINT\", \"Please set the endpoint\")\n",
    "\n",
    "# This is the name of the deployment specified in the Azure portal\n",
    "azure_deployment = os.getenv(\"AZUREAI_OPENAI_CHAT_DEPLOYMENT_NAME\", \"Please set the deployment name\")\n",
    "\n",
    "# This is the deployed API version, such as 2024-02-15-preview\n",
    "azure_api_version = os.getenv(\"AZUREAI_OPENAI_CHAT_API_VERSION\", \"Please set the API version\")\n",
    "\n",
    "# The environment variable should be set to the API key from the Azure AI playground:\n",
    "# api_key=os.getenv(\"AZUREAI_CHAT_KEY\", \"Please set API key\")\n",
    "\n",
    "# Alternatively, we can use Entra authentication\n",
    "token_provider = get_bearer_token_provider(\n",
    "     DefaultAzureCredential(),\n",
    "     \"https://cognitiveservices.azure.com/.default\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3203ea1-00b7-4cc6-9b92-5167d453a790",
   "metadata": {},
   "source": [
    "We can now construct the `guidance` model object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "29187b26-93bc-459f-bb9d-ac22feef473a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from guidance import models, gen\n",
    "from guidance.models import create_azure_openai_model\n",
    "\n",
    "azureai_model = create_azure_openai_model(\n",
    "    model_name=model,\n",
    "    azure_deployment=azure_deployment,\n",
    "    azure_endpoint=azure_endpoint,\n",
    "    api_version=azure_api_version,\n",
    "    # For authentication, use either\n",
    "    # api_key=api_key\n",
    "    # or\n",
    "    azure_ad_token_provider=token_provider,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f95b1ef-402a-4b96-b0b0-a0bd007d582b",
   "metadata": {},
   "source": [
    "We can use the model as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4b03a844-0c4c-446e-91ed-408494114d3a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51bdb8322c0141e2beeff23b44038e57",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "StitchWidget(initial_height='auto', initial_width='100%', srcdoc='<!doctype html>\\n<html lang=\"en\">\\n<head>\\n …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from guidance import system, user, assistant\n",
    "\n",
    "with system():\n",
    "    lm = azureai_model + \"You are a helpful assistant.\"\n",
    "    \n",
    "with user():\n",
    "    lm += \"What is the meaning of life?\"\n",
    "\n",
    "with assistant():\n",
    "    lm += gen(\"response\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "da5ea477",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "time.sleep(call_delay_secs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fa274c5-01c6-46f9-a024-b8b44e16ce2c",
   "metadata": {},
   "source": [
    "AOAI models also support constrained generation using JSON:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e8c9d8ff-99e2-4806-8ff5-7f57491d7361",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "80096e94cb064664a78cc159bece59b0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "StitchWidget(initial_height='auto', initial_width='100%', srcdoc='<!doctype html>\\n<html lang=\"en\">\\n<head>\\n …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"name\": \"Whiskers\",\n",
      "    \"age\": 3,\n",
      "    \"colour\": [\n",
      "        255,\n",
      "        200,\n",
      "        150\n",
      "    ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "from guidance import json as gen_json\n",
    "\n",
    "cat_schema = {\n",
    "    \"type\": \"object\",\n",
    "    \"properties\": {\n",
    "        \"name\": {\"type\": \"string\", \"minLength\": 4},\n",
    "        \"age\": {\"type\": \"integer\", \"minimum\": 0, \"maximum\": 20},\n",
    "        \"colour\": {\n",
    "            \"type\": \"array\",\n",
    "            \"items\": {\"type\": \"integer\", \"minimum\": 0, \"maximum\": 255},\n",
    "            \"minItems\": 3,\n",
    "            \"maxItems\": 3,\n",
    "        },\n",
    "    },\n",
    "    \"required\": [\"name\", \"age\", \"colour\"],\n",
    "    \"additionalProperties\": False,\n",
    "}\n",
    "\n",
    "with system():\n",
    "    lm = azureai_model + \"You are an expert in the ancient lore of cats\"\n",
    "\n",
    "with user():\n",
    "    lm += \"Create a simple description of a cat in JSON, including the name, age & colour\"\n",
    "\n",
    "with assistant():\n",
    "    lm += gen_json(schema=cat_schema, name=\"my_cat_text\", temperature=1.0)\n",
    "\n",
    "\n",
    "my_cat = json.loads(lm[\"my_cat_text\"])\n",
    "\n",
    "print(json.dumps(my_cat, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d5266819",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "time.sleep(call_delay_secs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c94fb3a-df2e-45e8-9fe5-134a0868dbf1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
